import socket
import gzip
import json
import io
import sys
from kafka import KafkaProducer

if len(sys.argv) != 5:
  print("Usage: script.py <udp_ip> <udp_port> <kafka_broker> <kafka_topic>")
  sys.exit(1)

UDP_IP = sys.argv[1]
UDP_PORT = int(sys.argv[2])
KAFKA_BROKER = sys.argv[3]
KAFKA_TOPIC = sys.argv[4]

sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
sock.bind((UDP_IP, UDP_PORT))

data_dict = {}

fields_to_remove = [
  "_id", "host", "level", "version", "_forwarder", "timestamp", "short_message",
  "_gl2_remote_ip", "_gl2_remote_port", "_gl2_source_node", "_gl2_source_collector",
  "_gl2_source_input", "_gl2_message_id", "_timestamp", "_message", "_SYSTEM",
  "_BYTES", "_CARD", "_AGENT", "_message-type", "_LICENSE",
  "_RETRANS_PKTS", "_CARD_IP", "_NIC", "_SYNACK_PKTS", "_L4_PROTOCOL", "_TOTAL_CPU",
  "_DISK_RATIO", "_RST_PKTS", "_PROBE_VER", "_RSP_BODY", "_PAGELOAD",
  "_SRC_NET_DELAY", "_DST_NET_DELAY", "_NET_DELAY", "_HTTP_RESPONSE", "_SQL_RESPONSE"
]

def process_url(url, domain):
  param = ''
  if '?' in url:
    split_list = url.split('?')
    if len(split_list) > 0: url = split_list[0]
    if len(split_list) > 1: param = split_list[1]
  if domain:
    url = f"http://{domain}{url}"
  return url, param

def process_header(header):
  user_agent = cookie = ''
  lines = header.split('\r\n')
  for line in lines:
    if line.startswith('User-Agent:'):
      user_agent = line.split(':', 1)[1].strip()
    elif line.startswith('Cookie:'):
      cookie = line.split(':', 1)[1].strip()
  
  return user_agent, cookie

def trim_body(body):
  return body[:100] + body[-100:] if len(body) > 200 else body

print(f"Listening on UDP port {UDP_PORT}")

producer = KafkaProducer(bootstrap_servers=[KAFKA_BROKER],
                         value_serializer=lambda v: json.dumps(v).encode('utf-8'))

while True:
  data, addr = sock.recvfrom(10240)

  try:
    with gzip.GzipFile(fileobj=io.BytesIO(data)) as f:
      decompressed_data = f.read()
    json_data = json.loads(decompressed_data)

    cleaned_data = {}
    cleaned_data['REQ_SIZE'] = 0
    cleaned_data['RSP_SIZE'] = 0
    if '_REQ_HEADER' in json_data: cleaned_data['REQ_SIZE'] += len(json_data['_REQ_HEADER'])
    if '_REQ_BODY'   in json_data: cleaned_data['REQ_SIZE'] += len(json_data['_REQ_BODY'])
    if '_RSP_HEADER' in json_data: cleaned_data['RSP_SIZE'] += len(json_data['_RSP_HEADER'])
    if '_RSP_BODY'   in json_data: cleaned_data['RSP_SIZE'] += len(json_data['_RSP_BODY'])

    if '_REQ_HEADER' in json_data: 
      user_agent, cookie = process_header(json_data['_REQ_HEADER'])
      cleaned_data['COOKIE'] = cookie
      cleaned_data['USER_AGENT'] = user_agent

    for key, value in json_data.items():
      new_key = key.lstrip('_')
      if key in fields_to_remove: continue
      if new_key == "URL": 
        path, param = process_url(value, json_data.get("_DOMAIN", ""))
        cleaned_data[new_key] = path
        cleaned_data['PARAM'] = param
      elif new_key in ["RSP_BODY", "REQ_BODY"]: 
        cleaned_data[new_key] = trim_body(value)
      else:
        cleaned_data[new_key] = value

    if "REQ_HEADER" not in cleaned_data: cleaned_data["REQ_HEADER"] = "" 
    if "RSP_HEADER" not in cleaned_data: cleaned_data["RSP_HEADER"] = "" 
    if "REQ_BODY"   not in cleaned_data: cleaned_data["REQ_BODY"]   = "" 
    if "RSP_BODY"   not in cleaned_data: cleaned_data["RSP_BODY"]   = "" 
    if "PARAM"      not in cleaned_data: cleaned_data["PARAM"]      = "" 
    if "COOKIE"     not in cleaned_data: cleaned_data["COOKIE"]     = "" 
    if "USER_AGENT" not in cleaned_data: cleaned_data["USER_AGENT"] = "" 

    l7 = cleaned_data["L7_PROTOCOL"]

    if len(cleaned_data) != 22 and l7 == 'http': 
      print(cleaned_data)

    if l7 != 'http' and l7 != 'dns':
      cleaned_data["DB_SERVER_NAME"] = l7
      cleaned_data["L7_PROTOCOL"] = "DB_SERVER"

    data_dict[addr] = cleaned_data
    producer.send(KAFKA_TOPIC, cleaned_data)
    producer.flush()
  except (json.JSONDecodeError) as e:
    pass

